import os, sys
import random
from PIL import Image
import base64
from io import BytesIO

import torch 
from torchvision import transforms

from CLIP import clip
from model import Generator

text_templates = [
    'a photo of {} style.',
    'a drawing of {} style.',
    'style of {} in a photo.',
    'the photo describes a {} style.',
    '{} style.',
    'the drawing is style of {}.',
    'the photo is {} style.',
    '{} drawing.'
]
CLIP_PATH = os.path.dirname(__file__) + '/exp/art_v2/models/200000-clip.ckpt'
G_PATH = os.path.dirname(__file__) + '/exp/art_v2/models/200000-G.ckpt'



class StyleTransfer:
    def __init__(self, clip_path=CLIP_PATH, generator_path=G_PATH, clip_model_name="RN50",
                emb_dim=1024, g_conv_dim=64, g_repeat_num=6, device=torch.device('cuda:0')):
        self.device = device
        self.encoder, _ = clip.load(clip_model_name, device, jit=False)
        self.encoder = self.encoder.float()

        self.G = Generator(emb_dim, g_conv_dim, g_repeat_num)
        self.encoder.load_state_dict(torch.load(clip_path, map_location=lambda storage, loc: storage))
        self.G.load_state_dict(torch.load(generator_path, map_location=lambda storage, loc: storage))

        self.encoder.to(device)
        self.G.to(device)
        self.encoder.eval()
        self.G.eval()

        self.trans = transforms.Compose(
            [transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
            ]
        )

    def denorm(self, x):
        """Convert the range from [-1, 1] to [0, 1]."""
        # out = (x + 1) / 2
        std = torch.FloatTensor((0.26862954, 0.26130258, 0.27577711), device=x.device).view(x.size(0), 3, 1, 1)
        mean = torch.FloatTensor((0.48145466, 0.4578275, 0.40821073), device=x.device).view(x.size(0), 3, 1, 1)
        out = std * x + mean
        return out.clamp_(0, 1)

    def forward_text(self, content_str: str, style_text_str: str):
        with torch.no_grad():
            content_bin = base64.b64decode(content_str)
            content_img = Image.open(BytesIO(content_bin))
            width, height = content_img.size
            c_img = self.trans(content_img).unsqueeze(0).to(self.device)

            format_ = random.choice(text_templates)
            style_text_str = format_.format(style_text_str)
            text_token = clip.tokenize(style_text_str).to(self.device)
            style_code = self.encoder.encode_text(text_token)
            res_img = self.G(c_img, style_code).squeeze(0)

            res_img = self.denorm(res_img.data.cpu())
            img = res_img.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
            im = Image.fromarray(img)
            im.resize((width, height))
        img_buffer = BytesIO()
        im.save(img_buffer, format='JPEG')
        byte_data = img_buffer.getvalue()
        return base64.b64encode(byte_data)

    def forward_image(self, content_str: str, style_img_str: str):
        with torch.no_grad():
            content_bin = base64.b64decode(content_str)
            content_img = Image.open(BytesIO(content_bin))
            width, height = content_img.size
            c_img = self.trans(content_img).unsqueeze(0).to(self.device)

            style_bin = base64.b64decode(style_img_str)
            style_img = Image.open(BytesIO(style_bin))
            s_img = self.trans(style_img).unsqueeze(0).to(self.device)

            style_code = self.encoder.encode_image(s_img)
            res_img = self.G(c_img, style_code).squeeze(0)

            res_img = self.denorm(res_img.data.cpu())
            img = res_img.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
            im = Image.fromarray(img)
            im.resize((width, height))
        img_buffer = BytesIO()
        im.save(img_buffer, format='JPEG')
        byte_data = img_buffer.getvalue()
        return base64.b64encode(byte_data)


if __name__ == '__main__':
    cbin = open('test/content.jpg', 'rb').read()
    cimg_str = base64.b64encode(cbin)

    sbin = open('test/style.jpg', 'rb').read()
    simg_str = base64.b64encode(sbin)

    stext = "Post Impressionism"
    
    trans = StyleTransfer()

    # use style image
    simg_res = trans.forward_image(cimg_str, simg_str)
    simg_bin = base64.b64decode(simg_res)
    simg = Image.open(BytesIO(simg_bin))
    simg.save("test/result_simg.jpg")

    # use style text
    simg_res = trans.forward_text(cimg_str, stext)
    simg_bin = base64.b64decode(simg_res)
    simg = Image.open(BytesIO(simg_bin))
    simg.save("test/result_stext.jpg")

    

